'''
This file Generate Train Set with format Train, Label
'''

import channelmodel as channel
import numpy as np
from scipy.io import savemat, loadmat
from Function import BSantennasNum, RISunitsNum, UEantenna, SNR, Carrierfrequency, Bandwidth, Pi, lamb, UnitsInGroup
from tqdm import tqdm
from multiprocessing import Pool


def TrainSetGenerate(snr = SNR):

    ChannelMatrix = loadmat('Data/ChannelMatrix.mat')
    BS2RIS = ChannelMatrix['BS2RIS']
    RIS2UE = ChannelMatrix['RIS2UE']

    # This is true Channel Matrix, label
    BS2UE = np.zeros((5000,16,64), dtype=complex)
    for j in range(5000):
        BS2UE[j] = BS2RIS[j] @ np.diag(RIS2UE[j])

    # This is caculated Channel Matrix, feature
    CALCUH = np.zeros((5000, BSantennasNum, int(RISunitsNum / UnitsInGroup)), dtype=complex)
    #---------------------------------------------------------------------------------------------------------------------------#
    
    # Train Date Generate and Write
    AllOmega = np.zeros((BSantennasNum, BSantennasNum), dtype=complex)
    AllPhi = np.zeros((int(RISunitsNum / UnitsInGroup), int(RISunitsNum / UnitsInGroup)), dtype=complex)
    AllRx = np.zeros((int(BSantennasNum * RISunitsNum / UnitsInGroup), 1), dtype=complex)

    for i in tqdm(range(5000)):
        # Fix the Phi
        for j in range(int(RISunitsNum / UnitsInGroup)):
            # TODO:could be simplified
            theta_Phi = np.zeros(int(RISunitsNum / UnitsInGroup))
            theta_Phi = [np.exp(- Pi / 6 + Pi / 3 / (RISunitsNum / UnitsInGroup) * j)
                        for j in range(int(RISunitsNum / UnitsInGroup))]
            Phi = RISReflectionMatrix(theta_Phi)
            AllPhi[j,:] = Phi

            # Change the Omega
            for k in range(BSantennasNum):
                # TODO:could be simplified
                theta_Omega = - Pi / 6 + Pi / 3 * (BSantennasNum) * k
                Omega = BeamformingMatrix(BSantennasNum, theta_Omega)
                tx = 1. / np.sqrt(BSantennasNum) * np.array(Omega)
                UPhi = np.kron(Phi, np.array([1,1]))
                rx = tx @ BS2RIS[i] * np.array(UPhi) @ RIS2UE[i] + \
                        np.random.normal(loc=0.0, scale=np.power(10,-(snr / 20)))
                AllOmega[k,:] = Omega
                AllRx[int(j * BSantennasNum) + k, 0] = rx

        W = np.kron(AllPhi, AllOmega)
        vecH = np.linalg.pinv(W) @ AllRx
        CALCUH[i] = vecH.reshape((int(RISunitsNum / UnitsInGroup), BSantennasNum)).T

    savemat('Data/NetworkData.mat', {'Label': BS2UE, 'feature': CALCUH})

        



def BeamformingMatrix(antennaNum, theta, d = 0.5 * lamb):
    omega = np.zeros((antennaNum, ), dtype=complex)
    for i in range(antennaNum):
        omega[i] = np.exp(1j * 2 * Pi * d * np.sin(theta) * i / lamb)

    return list(omega)
    
def RISReflectionMatrix(thetas):
    # theta must be a list of complex
    # return np.exp(1j * theta)
    return [np.exp(1j * theta) for theta in thetas]


if __name__ == '__main__':
        TrainSetGenerate(snr = 10)